{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test dynamic env class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import time, sys\n",
    "sys.path.append('..')\n",
    "import numpy as np\n",
    "from importlib import reload\n",
    "import matplotlib.pyplot as plt\n",
    "from pb_diff_envs.environment.static.maze2d_rand_wgrp import Maze2DRandRecGroupList\n",
    "import pb_diff_envs.environment.dynamic.dyn_rm2d_env as dre; reload(dre)\n",
    "from pb_diff_envs.environment.dynamic.dyn_rm2d_env import DynamicRandMaze2DEnv\n",
    "no_while = True\n",
    "\n",
    "# --- 2 ---\n",
    "wall_locations = np.array(\n",
    "    [np.array([1.0, 1.0]), np.array([4.2, 1.0]),])\n",
    "maze_size = np.array([5, 5])\n",
    "wall_hExts = np.array([[0.5, 0.5],[0.5, 0.5],])\n",
    "len_full_wtraj = 12\n",
    "num_sample_pts = 2\n",
    "gap_betw_wall = 0.15\n",
    "wp_seed = 100\n",
    "speed = 0.12 # 0.2 too long?\n",
    "\n",
    "robot_config = dict(maze_size=maze_size, min_to_wall_dist=0.01, collision_eps=0.05)\n",
    "env = DynamicRandMaze2DEnv(wall_locations, wall_hExts, \n",
    "                 len_full_wtraj, \n",
    "                 num_sample_pts,\n",
    "                 gap_betw_wall,\n",
    "                robot_config, wp_seed, speed, renderer_config={})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## simple vis of env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# env.robot.print_joint_states()\n",
    "# print(env.robot.item_id)\n",
    "# print(env.robot_id) # list\n",
    "print(env.robot.limits_low)\n",
    "print(env.robot.limits_high)\n",
    "print(env._get_joint_pose())\n",
    "print(f'{type(env._get_joint_pose())}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# env.resample_robot_start_pos()\n",
    "print(env.recWall_grp.get_walls_pos())\n",
    "print(env.recWall_grp.get_wtrajs_end_pos())\n",
    "print(env.recWall_grp.get_walls_pos_at_t(0))\n",
    "print(env.recWall_grp.is_walls_in_start_pos())\n",
    "pts = []\n",
    "for i in range(20000):\n",
    "    pts.append(env.get_robot_free_pose())\n",
    "pts = np.array(pts)\n",
    "# img = env.render()\n",
    "# img = env.render_1_traj(pts, env.recWall_grp.get_wtrajs(), img_type='png')\n",
    "from utils.maze2d_utils import plt_img\n",
    "# plt_img(img)\n",
    "# env.render_1_traj(pts)\n",
    "# env.recWall_grp.set_walls_at_time(0)\n",
    "plt.scatter(pts[:, 0], pts[:, 1])\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.xlim(-1, 6); plt.ylim(-1, 6)\n",
    "# plt_img(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample a problem/episode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from utils.utils import DotDict\n",
    "from tqdm import tqdm\n",
    "from utils.utils import seed_everything\n",
    "seed_everything(42) # 42\n",
    "# env.load(GUI=True) # False\n",
    "prev_pose = env.get_robot_free_pose()\n",
    "prev_pose = np.array([2.93287797, 3.33783522])\n",
    "env.robot.set_config(prev_pose)\n",
    "solutions = []\n",
    "wtrajs_list = []\n",
    "# n_traj = 25\n",
    "n_traj = 1\n",
    "for i in tqdm(range(n_traj)):\n",
    "    solution, d = env.sample_1_episode(prev_pose)\n",
    "\n",
    "    wtrajs = d['wtrajs'] # n_w, h, 2\n",
    "    wtrajs_list.append(wtrajs)\n",
    "\n",
    "    solution = np.array(solution)\n",
    "    prev_pose = solution[-1]\n",
    "    solutions.append(solution)\n",
    "\n",
    "## can concatenate all solutions\n",
    "# solutions = np.concatenate(solutions, axis=0)\n",
    "# wtrajs_list = np.concatenate(wtrajs_list, axis=1)\n",
    "# result = DotDict(solution=solutions)\n",
    "# print(solutions.shape) \n",
    "# print(len(result.solution))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "savepath = './test_sipp_dyn_maze2d.gif'\n",
    "sol_i = solutions[0]\n",
    "wtrajs_i = wtrajs_list[0]\n",
    "# gifs = env.render_1_traj(solutions, savepath=savepath, wtrajs=wtrajs_list) ## visualize all if concat\n",
    "gifs = env.render_1_traj(sol_i, savepath=savepath, wtrajs=wtrajs_i)\n",
    "savepath = glob(f\"{savepath.replace('.gif','')}**\")[0] ## retrieve the gif path\n",
    "###\n",
    "print(len(gifs), gifs[0].shape)\n",
    "from IPython.display import HTML\n",
    "import base64\n",
    "b64 = base64.b64encode(open(savepath, 'rb').read()).decode('ascii')\n",
    "display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(solution.shape) # (, 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "c1daed7e4d34e14fbb9bb260207f2e4074aea67593e999acdc484a19561ddc79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
